###########################################################
# © 2018-2019 Daniel 'grindhold' Brendle
#
# This file is part of elrotador.
#
# elrotador is free software: you can redistribute it and/or 
# modify it under the terms of the GNU Lesser General Public License 
# as published by the Free Software Foundation, either 
# version 3 of the License, or (at your option) any later 
# version.
#
# elrotador is distributed in the hope that it will be 
# useful, but WITHOUT ANY WARRANTY; without even the implied 
# warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR 
# PURPOSE. See the GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public 
# License along with elrotador. 
# If not, see http://www.gnu.org/licenses/.
###########################################################

import numpy as np
from math import *

import xml.etree.ElementTree


class Atom(object):
    def __init__(self, pos=None, typ=None):
        self.position = np.matrix([0,0,0])
        if pos is not None:
            self.position = pos
        self.type = "H"
        if typ is not None:
            self.type = typ

    def rotate(self, axis_origin, axis_direction, angle):
        """
        rotate this atom along a given axis specified by axis_origin and
        axis_direction
        """
        a = angle
        d = axis_direction

        norm = np.linalg.norm(d)
        if norm != 0:
            d = d/norm

        d = d[0]

        p = self.position - axis_origin
        rot = np.matrix([
            [
                (d[0]**2)*(1-cos(a))+cos(a),
                (d[0]*d[1])*(1-cos(a))-d[2]*sin(a), 
                (d[0]*d[2])*(1-cos(a))+d[1]*sin(a)
            ],[
                (d[1]*d[0])*(1-cos(a))+d[2]*sin(a),
                (d[1]**2)*(1-cos(a))+cos(a),
                (d[1]*d[2])*(1-cos(a))-d[0]*sin(a)
            ],[
                (d[2]*d[0])*(1-cos(a))-d[1]*sin(a),
                (d[2]*d[1])*(1-cos(a))+d[0]*sin(a),
                (d[2]**2)*(1-cos(a))+cos(a)
            ]
        ])
        p *= rot
        self.position = p + axis_origin

class Molecule(object):
    def __init__(self):
        self.atoms = []
        self.bonds = np.matrix([])
        self.charge = 0
        self.multiplicity = 1

    @classmethod
    def load_from_cml(cls, cml):
        o = cls()
        e = xml.etree.ElementTree.parse(cml).getroot()
        for el in e.getchildren():
            if el.tag == "atomArray":
                o.atoms = [None for x in range(len(el.getchildren()))]
                for atom in el.getchildren():
                    try:
                        ident = int(atom.attrib["id"].strip("a"))-1
                    except ValueError:
                        print("Remember implementing this for letters other than a if this exception drops")
                        continue

                    pos = np.matrix([
                        float(atom.attrib["x3"]),
                        float(atom.attrib["y3"]),
                        float(atom.attrib["z3"])
                    ])
                    o.atoms[ident] = Atom(pos, atom.attrib["elementType"])

        for el in e.getchildren():
            if el.tag == "bondArray":
                o.bonds = [[0 for _ in range(len(o.atoms))] for _ in range(len(o.atoms))]
                for bond in el.getchildren():
                    try:
                        atoms = [int(a.strip("a"))-1 for a in bond.attrib["atomRefs2"].split(" ")]
                    except ValueError:
                        print("Remember implementing this for letters other than a if this exception drops")
                        continue
                    if len(atoms) != 2:
                        raise Exception("A bond concerned more than 2 atoms")

                    o.bonds[atoms[0]][atoms[1]] = o.bonds[atoms[1]][atoms[0]] = int(bond.attrib["order"])
        return o

    def get_dihedral_wing(self, blind_side, active_side):
        if type(blind_side) != Atom or type(active_side) != Atom:
            raise TypeError("Inputs must be Atoms")

        if blind_side not in self.atoms or active_side not in self.atoms:
            raise ValueError("Atoms must be contained in this molecule")

        self.visited = set([blind_side])
        self._walk_wing(active_side)
        self.visited.remove(blind_side)
        return self.visited
    
    def _walk_wing(self, atom):
        self.visited.add(atom)
        atom_id = self.atoms.index(atom)
        neighbor_atoms = [self.atoms[idx] for idx, order in enumerate(self.bonds[atom_id]) if order > 0]
        for neighbor_atom in neighbor_atoms:
            if neighbor_atom not in self.visited:
                self._walk_wing(neighbor_atom)
            
        
    def rotate(self, dihedral, angle):
        blind_side = self.atoms[dihedral[1]-1]
        active_side = self.atoms[dihedral[2]-1]

        for atom in self.get_dihedral_wing(blind_side, active_side):
            atom.rotate(np.array(blind_side.position),
                        np.array((blind_side.position-active_side.position)*-1),
                        angle)

    def length(self, atoms, length):
        blind_side = self.atoms[atoms[0]-1]
        active_side = self.atoms[atoms[1]-1]

        d = blind_side.position - active_side.position
        norm = np.linalg.norm(d)
        if norm != 0:
            d /= norm

        for atom in self.get_dihedral_wing(blind_side, active_side):
            atom.position -= length * d

    def get_xyz(self):
        r = ""
        r += "    {}\n".format(len(self.atoms))
        r += "This is a jolly comment\n"
        for atom in self.atoms:
            r += "    {}    {}    {}    {}\n".format(
                            atom.type, 
                            atom.position.item(0),
                            atom.position.item(1),
                            atom.position.item(2)
            )
        return r

    def get_psi4(self):
        r = "units angstrom\n"
        r += "{} {}\n".format(self.charge, self.multiplicity)
        for atom in self.atoms:
            r += "    {}    {}    {}    {}\n".format(
                            atom.type, 
                            atom.position.item(0),
                            atom.position.item(1),
                            atom.position.item(2)
            )
        return r
